"""SQL解析模块 - 使用sqlglot解析CREATE TABLE语句"""

import re
import warnings
import sys
import io
import contextlib
from typing import List, Optional
import sqlglot
from sqlglot.expressions import Create, ColumnDef, DataType, Command

from .models import Table, Column


class SQLParser:
    """SQL解析器"""
    
    DIALECT_MAP = {
        'mysql': 'mysql',
        'postgres': 'postgres',
        'kingbase': 'postgres',  # 人大金仓基于PostgreSQL
    }
    
    def __init__(self, dialect: str = 'mysql'):
        """
        初始化解析器
        
        Args:
            dialect: 数据库类型，可选值：mysql, postgres, kingbase
        """
        self.dialect = dialect.lower()
        if self.dialect not in self.DIALECT_MAP:
            raise ValueError(f"不支持的数据库类型: {dialect}")
        self.sqlglot_dialect = self.DIALECT_MAP[self.dialect]
    
    def parse_file(self, file_path: str) -> List[Table]:
        """
        解析SQL文件，提取所有表结构
        
        Args:
            file_path: SQL文件路径
            
        Returns:
            表结构列表
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
        except FileNotFoundError:
            raise FileNotFoundError(f"文件不存在: {file_path}")
        except UnicodeDecodeError:
            raise ValueError(f"文件编码错误，请确保文件为UTF-8编码: {file_path}")
        
        return self.parse_sql(content)
    
    def parse_sql(self, sql_content: str) -> List[Table]:
        """
        解析SQL内容，提取所有表结构
        
        Args:
            sql_content: SQL内容字符串
            
        Returns:
            表结构列表
        """
        tables = []
        
        # 使用sqlglot解析SQL
        import re
        statements = []
        
        # 抑制sqlglot的警告（它会打印"contains unsupported syntax"等警告到stderr）
        stderr_buffer = io.StringIO()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            with contextlib.redirect_stderr(stderr_buffer):
                try:
                    # 尝试解析整个SQL文件
                    statements = sqlglot.parse(sql_content, dialect=self.sqlglot_dialect) or []
                    # 检查是否成功解析为CREATE TABLE语句（排除Command对象）
                    has_create_table = any(isinstance(stmt, Create) for stmt in statements)
                    if has_create_table:
                        # 如果成功解析，直接使用结果
                        pass
                    else:
                        # 如果没有CREATE TABLE语句（可能是Command对象），尝试MySQL方言
                        try:
                            statements_mysql = sqlglot.parse(sql_content, dialect='mysql') or []
                            if any(isinstance(stmt, Create) for stmt in statements_mysql):
                                statements = statements_mysql
                                has_create_table = True
                        except Exception:
                            pass
                except Exception:
                    # 如果整体解析失败，尝试MySQL方言作为fallback
                    try:
                        statements = sqlglot.parse(sql_content, dialect='mysql') or []
                        has_create_table = any(isinstance(stmt, Create) for stmt in statements)
                    except Exception:
                        statements = []
                        has_create_table = False
        
        # 如果解析结果为空或没有CREATE TABLE语句，使用fallback方法
        if not statements or not has_create_table:
            # 使用正则表达式提取CREATE TABLE语句
            # 匹配CREATE TABLE语句，包括多行和嵌套括号
            # 找到所有CREATE TABLE开始位置
            create_positions = []
            for match in re.finditer(r'CREATE\s+TABLE', sql_content, re.IGNORECASE):
                create_positions.append(match.start())
            
            # 清空之前的statements，使用fallback方法重新解析
            statements = []
            
            # 对每个CREATE TABLE，找到对应的结束分号
            for pos in create_positions:
                # 从CREATE TABLE开始，找到匹配的括号和结束分号
                paren_count = 0
                end_pos = pos
                i = pos
                found_open_paren = False
                while i < len(sql_content):
                    char = sql_content[i]
                    if char == '(':
                        paren_count += 1
                        found_open_paren = True
                    elif char == ')':
                        paren_count -= 1
                    elif char == ';' and paren_count == 0 and found_open_paren:
                        end_pos = i + 1
                        break
                    i += 1
                    # 安全限制：如果超过10000字符还没找到结束，可能有问题
                    if i - pos > 10000:
                        break
                
                if end_pos > pos:
                    sql_stmt = sql_content[pos:end_pos]
                    # 清理SQL语句：移除注释和多余空白
                    cleaned = re.sub(r'/\*.*?\*/', '', sql_stmt, flags=re.DOTALL)
                    cleaned = re.sub(r'--.*?\n', '\n', cleaned)
                    cleaned = cleaned.strip()
                    
                    if not cleaned:
                        continue
                    
                    parsed = None
                    # 抑制sqlglot的警告
                    with warnings.catch_warnings():
                        warnings.simplefilter("ignore")
                        with contextlib.redirect_stderr(io.StringIO()):
                            # 首先尝试使用指定的方言解析
                            try:
                                parsed = sqlglot.parse(cleaned, dialect=self.sqlglot_dialect)
                                # 检查是否成功解析为CREATE TABLE语句（排除Command对象）
                                if parsed and any(isinstance(stmt, Create) for stmt in parsed):
                                    statements.extend(parsed)
                                    continue
                            except Exception:
                                pass
                            
                            # 如果指定方言解析失败或没有CREATE TABLE语句，尝试MySQL方言作为fallback
                            # 这可以处理MySQL语法但指定了PostgreSQL/kingbase的情况
                            if not parsed or not any(isinstance(stmt, Create) for stmt in parsed):
                                try:
                                    parsed_mysql = sqlglot.parse(cleaned, dialect='mysql')
                                    if parsed_mysql and any(isinstance(stmt, Create) for stmt in parsed_mysql):
                                        statements.extend(parsed_mysql)
                                        continue
                                except Exception:
                                    pass
                    
                    # 如果还是失败，尝试更激进的清理
                    cleaned2 = cleaned
                    # 移除PostgreSQL特定的默认值表达式（如 nextval(...)）
                    cleaned2 = re.sub(r'\s+DEFAULT\s+nextval\([^)]+\)', '', cleaned2, flags=re.IGNORECASE)
                    # 移除列定义中的CHARACTER SET和COLLATE
                    cleaned2 = re.sub(r'\s+CHARACTER\s+SET\s+"[^"]*"\."[^"]*"', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+CHARACTER\s+SET\s+"[^"]*"', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+CHARACTER\s+SET\s+\w+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+COLLATE\s+"[^"]*"\."[^"]*"', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+COLLATE\s+"[^"]*"', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+COLLATE\s+\w+', '', cleaned2, flags=re.IGNORECASE)
                    # 移除表级别的选项
                    cleaned2 = re.sub(r'\s+ENGINE\s*=\s*\w+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+AUTO_INCREMENT\s*=\s*\d+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+DEFAULT\s+CHARSET\s*=\s*\w+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+COLLATE\s*=\s*\w+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+ROW_FORMAT\s*=\s*\w+', '', cleaned2, flags=re.IGNORECASE)
                    cleaned2 = re.sub(r'\s+TABLESPACE\s+\w+', '', cleaned2, flags=re.IGNORECASE)
                    # 移除USING BTREE等索引选项
                    cleaned2 = re.sub(r'\s+USING\s+\w+', '', cleaned2, flags=re.IGNORECASE)
                    # 移除NULLS FIRST/LAST
                    cleaned2 = re.sub(r'\s+NULLS\s+(FIRST|LAST)', '', cleaned2, flags=re.IGNORECASE)
                    
                    # 尝试用指定方言解析清理后的SQL
                    with warnings.catch_warnings():
                        warnings.simplefilter("ignore")
                        with contextlib.redirect_stderr(io.StringIO()):
                            try:
                                parsed = sqlglot.parse(cleaned2, dialect=self.sqlglot_dialect)
                                if parsed and any(isinstance(stmt, Create) for stmt in parsed):
                                    statements.extend(parsed)
                                    continue
                            except Exception:
                                pass
                            
                            # 最后尝试用MySQL方言解析清理后的SQL
                            try:
                                parsed = sqlglot.parse(cleaned2, dialect='mysql')
                                if parsed and any(isinstance(stmt, Create) for stmt in parsed):
                                    statements.extend(parsed)
                            except Exception:
                                continue
        
        if not statements:
            return tables
        
        for statement in statements:
            # 检查是否为CREATE TABLE语句
            # 对于Create对象，只要不是CREATE INDEX、CREATE VIEW等，就尝试解析为表
            if isinstance(statement, Create):
                kind_value = getattr(statement, 'kind', None)
                # sqlglot的kind可能是字符串"TABLE"、None或其他值
                # 如果kind不是明确的非表类型（如INDEX、VIEW），就尝试解析为表
                kind_str = str(kind_value).upper() if kind_value else ""
                if not kind_str or kind_str == "TABLE" or kind_str in ("", "NONE"):
                    table = self._parse_create_table(statement)
                    if table:
                        tables.append(table)
        
        # 解析 COMMENT ON COLUMN 语句（PostgreSQL 特有）
        self._parse_comment_on_column(sql_content, tables)
        
        return tables
    
    def _parse_create_table(self, create_stmt: Create) -> Optional[Table]:
        """解析CREATE TABLE语句"""
        # 获取表名
        # sqlglot中，表名在create_stmt.this.this.name中（this是Schema，this.this是Table）
        # 对于PostgreSQL的 "schema"."table" 格式，需要分别获取schema和table名
        table_name = None
        schema_name = None
        
        if hasattr(create_stmt, 'this') and create_stmt.this:
            schema = create_stmt.this
            # Schema对象可能有name属性（schema名）和this属性（Table对象）
            if hasattr(schema, 'name'):
                schema_name = str(schema.name).strip('`"\'')
            # Schema对象的this属性是Table对象
            if hasattr(schema, 'this'):
                table_obj = schema.this
                # Table对象的name属性包含表名
                if hasattr(table_obj, 'name'):
                    table_name = str(table_obj.name)
                elif hasattr(table_obj, 'this'):
                    table_name = str(table_obj.this)
                else:
                    table_name = str(table_obj)
            elif hasattr(schema, 'name') and not table_name:
                # 如果没有this属性，schema.name可能就是表名（无schema的情况）
                table_name = str(schema.name)
            else:
                table_name = str(schema)
        
        # 清理表名（移除反引号等）
        if table_name:
            table_name = table_name.strip('`"\'')
        
        # 如果有schema名，可以选择保留或移除
        # 这里选择保留schema前缀，格式为 schema.table
        if schema_name and table_name and schema_name != table_name:
            # 如果schema名和表名不同，组合它们
            full_table_name = f"{schema_name}.{table_name}"
        else:
            full_table_name = table_name
        
        if not full_table_name:
            return None
        
        # 使用完整表名（包括schema）创建Table对象
        table_name = full_table_name
        
        columns = []
        
        # 提取列定义
        # sqlglot中，列定义在create_stmt.this.expressions中（Schema对象的expressions属性）
        expressions = []
        if hasattr(create_stmt, 'this') and create_stmt.this:
            if hasattr(create_stmt.this, 'expressions'):
                expressions = create_stmt.this.expressions
            elif hasattr(create_stmt, 'expression') and create_stmt.expression:
                if hasattr(create_stmt.expression, 'expressions'):
                    expressions = create_stmt.expression.expressions
                elif isinstance(create_stmt.expression, list):
                    expressions = create_stmt.expression
                else:
                    expressions = [create_stmt.expression]
        
        for expr in expressions:
            if isinstance(expr, ColumnDef):
                column = self._parse_column(expr)
                if column:
                    columns.append(column)
        
        # 提取表注释（从表级别属性）
        table_comment = None
        # MySQL: COMMENT='xxx'
        # 尝试从原始SQL中提取表注释（备用方案）
        
        return Table(name=table_name, columns=columns, comment=table_comment)
    
    def _parse_column(self, column_def: ColumnDef) -> Optional[Column]:
        """解析列定义"""
        if not hasattr(column_def, 'this') or not column_def.this:
            return None
        
        # 获取列名
        column_name = None
        if hasattr(column_def.this, 'name'):
            column_name = column_def.this.name
        else:
            column_name = str(column_def.this)
        
        if not column_name:
            return None
        
        # 提取数据类型
        data_type = "VARCHAR"
        length = None
        
        if hasattr(column_def, 'kind') and column_def.kind:
            if isinstance(column_def.kind, DataType):
                # 获取类型名
                if hasattr(column_def.kind, 'this'):
                    data_type = str(column_def.kind.this).upper()
                else:
                    data_type = str(column_def.kind).upper()
                
                # 移除 "TYPE." 前缀（如果存在）
                if data_type.startswith('TYPE.'):
                    data_type = data_type[5:]  # 移除 "TYPE." 前缀
                
                # 提取长度/精度
                if hasattr(column_def.kind, 'expressions') and column_def.kind.expressions:
                    length_parts = []
                    for expr in column_def.kind.expressions:
                        length_parts.append(str(expr))
                    if length_parts:
                        length = ','.join(length_parts)
        
        # 提取约束和属性
        nullable = True
        default = None
        comment = None
        
        # 检查约束
        if hasattr(column_def, 'constraints') and column_def.constraints:
            for constraint in column_def.constraints:
                constraint_type = None
                constraint_kind = None
                if hasattr(constraint, 'kind'):
                    constraint_kind = constraint.kind
                    constraint_type = str(constraint_kind).upper()
                else:
                    constraint_type = str(constraint).upper()
                
                # COMMENT约束（MySQL）
                # sqlglot中COMMENT约束的kind是CommentColumnConstraint对象
                if constraint_type.startswith('COMMENT') or 'COMMENT' in constraint_type:
                    # CommentColumnConstraint的this属性包含注释值（Literal对象）
                    if hasattr(constraint_kind, 'this'):
                        comment = self._extract_string_value(constraint_kind.this)
                    elif hasattr(constraint, 'expression'):
                        comment = self._extract_string_value(constraint.expression)
                
                # NOT NULL约束
                elif constraint_type == 'NOT NULL' or 'NOT NULL' in constraint_type:
                    nullable = False
                
                # DEFAULT约束
                # sqlglot中DEFAULT约束的kind是DefaultColumnConstraint对象
                elif constraint_type == 'DEFAULT' or constraint_type.startswith('DEFAULT'):
                    # DefaultColumnConstraint的this属性包含默认值
                    if hasattr(constraint_kind, 'this'):
                        default = self._extract_default_value(constraint_kind.this)
                    elif hasattr(constraint, 'expression'):
                        default = self._extract_default_value(constraint.expression)
        
        return Column(
            name=column_name,
            type=data_type,
            length=length,
            nullable=nullable,
            default=default,
            comment=comment
        )
    
    def _extract_string_value(self, expr) -> Optional[str]:
        """提取字符串值"""
        if expr is None:
            return None
        
        # 尝试直接获取字符串值
        if hasattr(expr, 'this'):
            value = expr.this
            if isinstance(value, str):
                # 移除引号
                return value.strip('\'"')
            return str(value).strip('\'"')
        
        # 尝试从表达式中提取
        expr_str = str(expr)
        # 移除引号
        expr_str = expr_str.strip('\'"')
        return expr_str if expr_str else None
    
    def _extract_default_value(self, expr) -> Optional[str]:
        """提取默认值"""
        if expr is None:
            return None
        
        # 处理函数类型的默认值（如CURRENT_TIMESTAMP）
        expr_str = str(expr)
        expr_upper = expr_str.upper()
        if 'CURRENT_TIMESTAMP' in expr_upper or 'CURRENT_TIME' in expr_upper or 'CURRENT_DATE' in expr_upper or 'NOW()' in expr_upper:
            # 提取函数名（去掉括号）
            if 'CURRENT_TIMESTAMP' in expr_upper:
                return 'CURRENT_TIMESTAMP'
            elif 'CURRENT_TIME' in expr_upper:
                return 'CURRENT_TIME'
            elif 'CURRENT_DATE' in expr_upper:
                return 'CURRENT_DATE'
            elif 'NOW()' in expr_upper:
                return 'NOW()'
        
        # 对于字符串值，保留引号；对于其他值，直接转换为字符串
        if hasattr(expr, 'this'):
            value = expr.this
            if value is None:
                # 如果this为None，可能是函数类型，使用字符串表示
                return expr_str
            if isinstance(value, str):
                # 字符串值保留引号
                return f"'{value}'"
            # 数字、布尔等直接转字符串
            return str(value)
        
        # 直接转换为字符串
        return expr_str
    
    def _parse_comment_on_column(self, sql_content: str, tables: List[Table]):
        """
        解析 COMMENT ON COLUMN 语句，将备注信息关联到对应的列
        
        Args:
            sql_content: SQL内容字符串
            tables: 已解析的表列表
        """
        # 创建表名到Table对象的映射（支持schema.table格式）
        table_map = {}
        for table in tables:
            # 表名可能是 "schema.table" 或 "table" 格式
            table_map[table.name] = table
            # 如果表名包含schema，也支持直接用表名查找
            if '.' in table.name:
                _, table_only = table.name.rsplit('.', 1)
                table_map[table_only] = table
        
        # 使用正则表达式匹配 COMMENT ON COLUMN 语句
        # 格式: COMMENT ON COLUMN "schema"."table"."column" IS 'comment';
        # 支持两种格式：
        # 1. COMMENT ON COLUMN "schema"."table"."column" IS 'comment'
        # 2. COMMENT ON COLUMN "table"."column" IS 'comment'
        # 先尝试匹配三部分格式（schema.table.column）
        pattern1 = r'COMMENT\s+ON\s+COLUMN\s+"([^"]+)"\s*\.\s*"([^"]+)"\s*\.\s*"([^"]+)"\s+IS\s+[\'"]([^\'"]*)[\'"]'
        for match in re.finditer(pattern1, sql_content, re.IGNORECASE | re.MULTILINE):
            schema = match.group(1)
            table_name = match.group(2)
            column_name = match.group(3)
            comment = match.group(4)
            
            # 查找对应的表
            target_table = None
            # 优先使用完整表名（schema.table）
            full_table_name = f"{schema}.{table_name}"
            target_table = table_map.get(full_table_name)
            
            # 如果没找到，尝试只用表名
            if not target_table:
                target_table = table_map.get(table_name)
            
            # 如果还是没找到，尝试匹配所有表（处理大小写不敏感的情况）
            if not target_table:
                for table_key, table_obj in table_map.items():
                    if table_key.lower() == full_table_name.lower() or table_key.lower() == table_name.lower():
                        target_table = table_obj
                        break
            
            # 找到表后，查找对应的列并设置备注
            if target_table:
                for column in target_table.columns:
                    if column.name.lower() == column_name.lower():
                        column.comment = comment
                        break
        
        # 再尝试匹配两部分格式（table.column，无schema）
        pattern2 = r'COMMENT\s+ON\s+COLUMN\s+"([^"]+)"\s*\.\s*"([^"]+)"\s+IS\s+[\'"]([^\'"]*)[\'"]'
        for match in re.finditer(pattern2, sql_content, re.IGNORECASE | re.MULTILINE):
            table_name = match.group(1)
            column_name = match.group(2)
            comment = match.group(3)
            
            # 查找对应的表
            target_table = table_map.get(table_name)
            
            # 如果没找到，尝试匹配所有表（处理大小写不敏感的情况）
            if not target_table:
                for table_key, table_obj in table_map.items():
                    if table_key.lower() == table_name.lower():
                        target_table = table_obj
                        break
            
            # 找到表后，查找对应的列并设置备注
            if target_table:
                for column in target_table.columns:
                    if column.name.lower() == column_name.lower():
                        column.comment = comment
                        break

